import math
import numpy as np
import open3d as o3d
from vis_utils import Vis, pc_tf
import os
import csv
from tqdm import tqdm
import cv2
import argparse
from utils import topic2path
from my_io import copy_file, read_points, write_points

def check_path(data_path, topics):
    topic_uni = []
    for topic in topics:
        path = os.path.join(data_path, topic2path(topic))
        if not os.path.exists(path):
            continue
        if topic in topic_uni:
            continue
        topic_uni.append(topic)
    return topic_uni


def main(args):
    assert os.path.exists(args.data_path)
    frames = os.listdir(args.data_path)
    frames.sort(key=lambda x: x)
    for dir in frames:
        if dir[:2] != '__':# '__'
            dir = os.path.join(args.data_path, dir)
            if os.path.isdir(dir):
                main_worker(dir, args)


def main_worker(frames_path, args):
    print("开始同步数据...")

    lidar_topics = ["/pointcloud_lidar1", "/pointcloud_lidar2", "/pointcloud_lidar3", "/pointcloud_lidar4"]
    radar_topics = ["/pointcloud_radar3"]
    camera_topics = ["/camera71/compressed", "/camera72/compressed", "/camera73/compressed", "/camera74/compressed",
                     "/camera75/compressed", "/camera76/compressed",
                     "/camera77/compressed", "/camera78/compressed", "/camera79/compressed", "/camera80/compressed",
                     "/camera81/compressed", "/camera82/compressed"]
    # 去重
    lidar_topics = check_path(os.path.join(frames_path, "lidar"), lidar_topics)
    radar_topics = check_path(os.path.join(frames_path, "radar"), radar_topics)
    camera_topics = check_path(os.path.join(frames_path, "camera"), camera_topics)

    # 读取lidar和radar外参
    sensor_rts = dict()
    fs = cv2.FileStorage(os.path.join(frames_path, 'calib', 'sensor.yml'), cv2.FileStorage_READ)
    for topic in lidar_topics + radar_topics:
        sensor_id = topic[-6:]
        RT = fs.getNode(sensor_id).mat()
        sensor_rts[sensor_id] = RT

    datas_stamps = dict()
    sensor_topics = []
    sensor_types = []
    for i in range(len(lidar_topics)):
        sensor_types.append('lidar')
    sensor_topics = sensor_topics + lidar_topics
    for i in range(len(radar_topics)):
        sensor_types.append('radar')
    sensor_topics = sensor_topics + radar_topics
    for i in range(len(camera_topics)):
        sensor_types.append('camera')
    sensor_topics = sensor_topics + camera_topics

    for topic, type in zip(sensor_topics, sensor_types):
        data_stamps = []
        path = os.path.join(frames_path, type, topic2path(topic))
        if len(os.listdir(path)) <= 0:
            print("{}文件夹内部没有数据,移除该topic".format(topic))
            if type == "lidar":
                lidar_topics.remove(topic)
            elif type == "radar":
                radar_topics.remove(topic)
            elif type == "camera":
                camera_topics.remove(topic)
            else:
                print("类型不属于['lidar','radar','camera']")
                exit(-1)
            continue
        data_files = os.listdir(path)
        data_files.sort(key=lambda x: int(x.split('.')[0]))
        for data_file in data_files:
            data_stamp = int(str.split(data_file, '.')[0])
            data_stamps.append(data_stamp)
        datas_stamps[topic] = data_stamps

    data_speeds = dict()
    filename = os.path.join(frames_path, "localization", "localization.csv")
    with open(filename, "r") as csvfile:
        csvreader = csv.reader(csvfile)
        data_stamps = []
        # 遍历csvreader对象的每一行内容并输出
        is_header = True  # 第一行表头不要
        for row in csvreader:
            if is_header:
                is_header = False
                continue
            data_stamps.append(int(row[0]))
            data_speeds[int(row[0])] = math.sqrt(float(row[7]) * float(row[7]) + float(row[8]) * float(row[8]) + float(row[9]) * float(row[9]))
        datas_stamps["/localization_result"] = data_stamps

    print("各话题对应的帧数")
    for topic in datas_stamps:
        print("话题<{}>   帧数<{}>".format(topic, len(datas_stamps[topic])))

    # 以前激光为中心去同步其他的数据
    center_topic = "/pointcloud_lidar3"
    datas_pointer = dict()
    sync_datas_dict = dict()
    for topic in datas_stamps:
        datas_pointer[topic] = 0
        sync_datas_dict[topic] = []

    for cencer_stamp in tqdm(datas_stamps[center_topic], postfix="同步数据..."):
        # 遍历每个topic
        for topic in datas_pointer:
            for i in range(datas_pointer[topic], len(datas_stamps[topic]) - 1):
                if (cencer_stamp < datas_stamps[topic][i]):
                    sync_datas_dict[topic].append(datas_stamps[topic][i])
                    datas_pointer[topic] = i
                    break
                if (datas_stamps[topic][i] <= cencer_stamp and cencer_stamp <= datas_stamps[topic][i + 1]):
                    delta_left = cencer_stamp - datas_stamps[topic][i]
                    delta_right = datas_stamps[topic][i + 1] - cencer_stamp
                    if (delta_left <= delta_right):
                        sync_datas_dict[topic].append(datas_stamps[topic][i])
                    else:
                        sync_datas_dict[topic].append(datas_stamps[topic][i + 1])
                    datas_pointer[topic] = i
                    break

    # 同步后某些topic的data数据不足则补0
    sync_num = len(sync_datas_dict[center_topic])
    for topic in sync_datas_dict:
        if len(sync_datas_dict[topic]) > sync_num:
            exit(-1)
        sync_datas_dict[topic] = sync_datas_dict[topic] + [0] * (sync_num - len(sync_datas_dict[topic]))

    # 从下面开始list变成array
    # 检测同步数据质量 diff stamp最大值不得大于delta_stamp_max_thre
    delta_stamp_max = np.zeros(sync_num)
    for topic in sync_datas_dict:
        delta_stamp = np.abs(np.array(sync_datas_dict[center_topic]) - np.array(sync_datas_dict[topic]))
        is_larger = delta_stamp > delta_stamp_max
        delta_stamp_max = delta_stamp * is_larger + delta_stamp_max * (1 - is_larger)
    valid_data_idx = delta_stamp_max < args.delta_stamp_max_thre  # 同步质量达标的数据置true
    # 只保存达标的数据
    for topic in sync_datas_dict:
        sync_datas_dict[topic] = np.array(sync_datas_dict[topic])[valid_data_idx]

    # jump frames when high-speed or full size in low-speed
    keep_data_idx = []
    jump = args.jump
    for localization_stamp in sync_datas_dict['/localization_result']:
        speed = data_speeds[localization_stamp]
        if speed > args.jump_speed or jump == args.jump:
            keep_data_idx.append(True)
            jump = 1#从1到jump循环
        else:
            jump = jump + 1
            keep_data_idx.append(False)
    keep_data_idx = np.array(keep_data_idx)
    # 只保存达标的数据
    for topic in sync_datas_dict:
        sync_datas_dict[topic] = np.array(sync_datas_dict[topic])[keep_data_idx]
    valid_data_sum = keep_data_idx.sum()  # 同步质量达标的数据量

    if args.is_view:
        vis = Vis()

    # 准备save数据
    # 新建目录
    output_paths = frames_path.replace(frames_path.split('/')[-1], '__' + frames_path.split('/')[-1])  # 输出文件夹为 “__源文件夹”
    output_paths_lidar = os.path.join(output_paths, "lidar")
    output_paths_lidar_bin = os.path.join(output_paths, "lidar_bin")
    output_paths_radar = os.path.join(output_paths, "radar_bin")#此处给radar改了名
    output_paths_camera = os.path.join(output_paths, "camera")
    output_paths_calib = os.path.join(output_paths, "calib")
    # output_paths_label = os.path.join(output_paths, "label")
    output_paths_localization = os.path.join(output_paths, "localization")
    output_paths_cameras = []
    for image_topic in camera_topics:
        output_paths_cameras.append(os.path.join(output_paths_camera, topic2path(image_topic)))
    for output_path in [output_paths, output_paths_lidar, output_paths_lidar_bin,output_paths_radar, output_paths_camera,
                        output_paths_calib,output_paths_localization] + output_paths_cameras:
        if not os.path.exists(output_path):
            print("新建目录<{}>".format(output_path))
            os.makedirs(output_path, exist_ok=True)
    # 把源文件夹内的标定参数和定位信息拷贝到新文件夹
    copy_file(os.path.join(frames_path, 'calib', 'sensor.yml'), os.path.join(output_paths, 'calib'))
    copy_file(os.path.join(frames_path, 'localization', 'localization.csv'), os.path.join(output_paths, 'localization'))
    # 新建readme说明文件
    # 创建或打开文件
    csvfile = open(os.path.join(output_paths,  'README.csv'), mode='w', newline='')
    # 标题列表
    fieldnames = ['Item', 'Details']
    # 创建 DictWriter 对象
    write = csv.DictWriter(csvfile, fieldnames=fieldnames)
    # 写入表头
    write.writeheader()
    write.writerow({'Item': 'Bag_Name', 'Details': frames_path.split('/')[-1]})
    write.writerow({'Item': 'With_Intensity', 'Details': True})
    write.writerow({'Item': 'Description', 'Details': '<you can note something here>'})
    write.writerow({'Item': 'SyncParas', 'Details': 'jump:{} jump_speed:{} delta_stamp_max_thre:{}'.format(args.jump, args.jump_speed, args.delta_stamp_max_thre)})
    write.writerow({'Item': 'Frames_Size', 'Details': len(sync_datas_dict[center_topic])})
    write.writerow({'Item': 'Samples_Size', 'Details': "Waitting For Sampling"})
    write.writerow({'Item': 'Topics_Size', 'Details': len(sync_datas_dict)})
    write.writerow({'Item': 'Topics', 'Details': list(sync_datas_dict.keys())})
    write.writerow({'Item': 'Split(train:test:val)', 'Details': "Waitting For Split"})

    for i in tqdm(range(0, valid_data_sum), postfix="save数据..."):
        # 拼接多个激光并save
        lidar_all = np.empty(shape=(0, 4))
        for lidar_topic in lidar_topics:
            bin_path = os.path.join(frames_path, 'lidar', topic2path(lidar_topic),
                                    str(sync_datas_dict[lidar_topic][i]) + '.bin')
            pc = read_points(bin_path, dim=4)
            lidar_tf = sensor_rts[lidar_topic[-6:]]
            points_trans = pc_tf(pc, lidar_tf)
            lidar_all = np.concatenate((lidar_all, points_trans))
        if args.is_save:
            sync_name = str(sync_datas_dict[center_topic][i])  # 以同步中心topic的时间戳为准，对其他topic进行重命名
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(lidar_all[:, :3])
            o3d.io.write_point_cloud(os.path.join(output_paths_lidar, sync_name + ".pcd"), pcd, write_ascii=True)#save xyz.pcd
            write_points(lidar_all, os.path.join(output_paths_lidar_bin, sync_name + ".bin"))#save xyzi.bin

        # 拼接多个radar并save
        radar_all = np.empty(shape=(0, 5))# x y z r_speed rcs 不建议使用r_speed
        for radar_topic in radar_topics:
            bin_path = os.path.join(frames_path, 'radar', topic2path(radar_topic),
                                    str(sync_datas_dict[radar_topic][i]) + '.bin')
            pc = read_points(bin_path, dim=5)
            radar_tf = sensor_rts[radar_topic[-6:]]
            points_trans = pc_tf(pc, radar_tf)
            radar_all = np.concatenate((radar_all, points_trans))
        if args.is_save:
            sync_name = str(sync_datas_dict[center_topic][i])  # 以同步中心topic的时间戳为准，对其他topic进行重命名
            write_points(radar_all,os.path.join(output_paths_radar, sync_name + ".bin"))


            # pcd = o3d.io.read_point_cloud(pcd_path)
            # pc = np.asarray(pcd.points)
            # radar_tf = sensor_rts[radar_topic[-6:]]
            # points_trans = pc_tf(pc, radar_tf)
            # radar_all = np.concatenate((radar_all, points_trans))
            #
            # if args.is_save:
            #     sync_name = str(sync_datas_dict[center_topic][i])  # 以同步中心topic的时间戳为准，对其他topic进行重命名
            #     pcd = o3d.geometry.PointCloud()
            #     pcd.points = o3d.utility.Vector3dVector(radar_all[:, :3])
            #     o3d.io.write_point_cloud(os.path.join(output_paths_radar, sync_name + ".pcd"), pcd, write_ascii=True)

        # save image
        images = []
        for image_topic in camera_topics:
            img_path = os.path.join(frames_path, 'camera', topic2path(image_topic),
                                    str(sync_datas_dict[image_topic][i]) + '.jpg')
            image = cv2.imread(img_path)
            images.append(image)
            if args.is_save:
                cv2.imwrite(os.path.join(output_paths, 'camera', topic2path(image_topic), sync_name + ".jpg"), image)

        if args.is_view:
            # Note all frames must be of the same size
            # image = cv2.resize(image, (0, 0), None, .4, .4)
            # grey = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            # grey_3_channel = cv2.cvtColor(grey, cv2.COLOR_GRAY2BGR)
            #
            # numpy_horizontal_row1 = np.hstack((image, grey_3_channel))
            # numpy_horizontal_row2 = np.hstack((grey_3_channel, image))
            # combined_images = np.concatenate((numpy_horizontal_row1, numpy_horizontal_row2), axis=0)
            cv2.imshow('Image panel', images[0])
            cv2.waitKey(1)

            # o3d.visualization.draw_geometries([img], window_name="Open3D显示图像",
            #                                   width=1024, height=768,
            #                                   left=50, top=50,
            #                                   mesh_show_back_face=False)  # 显示图片
            # if args.is_view:
            #     vis.show_img(img)

        if args.is_view:
            vis.show_pc(lidar_all, 'lidar', color=[0.2, 0.3, 0.4])
            vis.show_pc(radar_all, 'radar', color=[1, 0.706, 0.1])
            boxes = np.array([[0, 0, 0, 1, 1, 1, 0], [5, 5, 5, 2, 1, 1, 0.3]
                              ])
            vis.show_box(boxes)
            vis.show()

    print("同步完毕")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--data-path', default="/home/shn/Project/parse_result",
                        help='your data root for kitti')
    parser.add_argument('--is-view', default=False, action='store_false')  # 是否可视化
    parser.add_argument('--is-save', default=True, action='store_true')  # 是否保存
    parser.add_argument('--jump', type=int, default=1,help='get one frame among {jump} frames in lidar3') # 每次累计跳{jump}帧
    parser.add_argument('--jump_speed', type=float, default=3.0,help='') # 速度高于{jump_speed}开始跳帧
    parser.add_argument('--delta_stamp_max_thre', type=float, default=25.0,help='') # 同步的时间阈值
    args = parser.parse_args()

    main(args)
